// openni_tracker.cpp

#include <ros/ros.h>
#include <ros/package.h>
#include <tf/transform_broadcaster.h>
#include <kdl/frames.hpp>

#include <XnOpenNI.h>
#include <XnCodecIDs.h>
#include <XnCppWrapper.h>

#include <human_tracking/BodyArray.h>
#include <human_tracking/IndicesArray.h>

using std::string;

#define CHECK_RC(nRetVal, what)                                         \
  if (nRetVal != XN_STATUS_OK)                                          \
    {                                                                   \
      ROS_ERROR("%s failed: %s", what, xnGetStatusString(nRetVal));     \
      return nRetVal;                                                   \
    }
#define CHECK_RCN(nRetVal, what)                                        \
  if (nRetVal != XN_STATUS_OK)                                          \
    {                                                                   \
      ROS_ERROR("%s failed: %s", what, xnGetStatusString(nRetVal));     \
      exit(0);                                                          \
    }

class UserTracker {
  ros::NodeHandle nh_;
  ros::NodeHandle pnh_;

  bool register_depth_;
  bool save_mode_;

  std::string frame_id_;
  int queue_size_;

  ros::Publisher pub_bodies_;
  ros::Publisher pub_indices_;

public:
  bool use_calibfile_;

  xn::Context        g_Context;
  std::string calib_fname;
  xn::DepthGenerator g_DepthGenerator;
  xn::ImageGenerator g_ImageGenerator;
  xn::UserGenerator  g_UserGenerator;

  UserTracker () : nh_(), pnh_("~"), queue_size_(3) {
    calib_fname = ros::package::getPath("human_tracking") + "/user_calib.bin";
    register_depth_ = true;

    string configFilename = ros::package::getPath("human_tracking") + "/human_tracking.xml";
    XnStatus nRetVal = g_Context.InitFromXmlFile(configFilename.c_str());
    CHECK_RCN(nRetVal, "InitFromXml");

    nRetVal = g_Context.FindExistingNode(XN_NODE_TYPE_DEPTH, g_DepthGenerator);
    CHECK_RCN(nRetVal, "Find depth generator");

    nRetVal = g_Context.FindExistingNode(XN_NODE_TYPE_IMAGE, g_ImageGenerator);
    CHECK_RCN(nRetVal, "Find Image generator");

    if ( register_depth_ ) {
      XnStatus status = g_DepthGenerator.GetAlternativeViewPointCap ().SetViewPoint (g_ImageGenerator);
      if (status != XN_STATUS_OK) {
        //printf("not registered\n");
      }
      bool ret = g_DepthGenerator.GetAlternativeViewPointCap ().IsViewPointAs (g_ImageGenerator);
      ROS_INFO("registered ret = %d\n", ret);
    }

    nRetVal = g_Context.FindExistingNode(XN_NODE_TYPE_USER, g_UserGenerator);
    if (nRetVal != XN_STATUS_OK) {
      nRetVal = g_UserGenerator.Create(g_Context);
      CHECK_RCN(nRetVal, "Find user generator");
    }

    //
    if (!g_UserGenerator.IsCapabilitySupported(XN_CAPABILITY_SKELETON)) {
      ROS_ERROR("Supplied user generator doesn't support skeleton");
      exit(0);
    }

    if ( register_depth_ ) {
      frame_id_ = "openni_rgb_frame";
    } else {
      frame_id_ = "openni_depth_frame";
    }
    pnh_.getParam("camera_frame_id", frame_id_);

    pnh_.param("use_calib_file", use_calibfile_, true);
    pnh_.param("save_calib_file", save_mode_, false);

    pub_bodies_ = pnh_.advertise<human_tracking::BodyArray> ("bodies", queue_size_);
    pub_indices_ = pnh_.advertise<human_tracking::IndicesArray> ("indices", queue_size_);
  }

  void publishTransformTF(XnUserID const& user, XnSkeletonJoint const& joint, string const& frame_id,
                          string const& child_frame_id, ros::Time &tm) {
    static tf::TransformBroadcaster br;

    XnSkeletonJointPosition joint_position;
    g_UserGenerator.GetSkeletonCap().GetSkeletonJointPosition(user, joint, joint_position);
    double x = -joint_position.position.X / 1000.0;
    double y = joint_position.position.Y / 1000.0;
    double z = joint_position.position.Z / 1000.0;

    XnSkeletonJointOrientation joint_orientation;
    g_UserGenerator.GetSkeletonCap().GetSkeletonJointOrientation(user, joint, joint_orientation);

    XnFloat* m = joint_orientation.orientation.elements;
    KDL::Rotation rotation(m[0], m[1], m[2],
                           m[3], m[4], m[5],
                           m[6], m[7], m[8]);
    double qx, qy, qz, qw;
    rotation.GetQuaternion(qx, qy, qz, qw);

    char child_frame_no[128];
    snprintf(child_frame_no, sizeof(child_frame_no), "%s_%d", child_frame_id.c_str(), user);

    tf::Transform transform;
    transform.setOrigin(tf::Vector3(x, y, z));
    transform.setRotation(tf::Quaternion(qx, -qy, -qz, qw));

    // #4994
    tf::Transform change_frame;
    change_frame.setOrigin(tf::Vector3(0, 0, 0));
    tf::Quaternion frame_rotation;
    frame_rotation.setEulerZYX(1.5708, 0, 1.5708);
    change_frame.setRotation(frame_rotation);

    transform = change_frame * transform;

    br.sendTransform(tf::StampedTransform(transform, tm, frame_id, child_frame_no));
  }

  void publishTransformsTF(const std::string& frame_id, ros::Time &tm) {
    XnUserID users[15];
    XnUInt16 users_count = 15;
    g_UserGenerator.GetUsers(users, users_count);

    for (int i = 0; i < users_count; ++i) {
      XnUserID user = users[i];
      if (!g_UserGenerator.GetSkeletonCap().IsTracking(user))
        continue;

      publishTransformTF(user, XN_SKEL_HEAD,           frame_id, "head", tm);
      publishTransformTF(user, XN_SKEL_NECK,           frame_id, "neck", tm);
      publishTransformTF(user, XN_SKEL_TORSO,          frame_id, "torso", tm);

      publishTransformTF(user, XN_SKEL_LEFT_SHOULDER,  frame_id, "left_shoulder", tm);
      publishTransformTF(user, XN_SKEL_LEFT_ELBOW,     frame_id, "left_elbow", tm);
      publishTransformTF(user, XN_SKEL_LEFT_HAND,      frame_id, "left_hand", tm);

      publishTransformTF(user, XN_SKEL_RIGHT_SHOULDER, frame_id, "right_shoulder", tm);
      publishTransformTF(user, XN_SKEL_RIGHT_ELBOW,    frame_id, "right_elbow", tm);
      publishTransformTF(user, XN_SKEL_RIGHT_HAND,     frame_id, "right_hand", tm);

      publishTransformTF(user, XN_SKEL_LEFT_HIP,       frame_id, "left_hip", tm);
      publishTransformTF(user, XN_SKEL_LEFT_KNEE,      frame_id, "left_knee", tm);
      publishTransformTF(user, XN_SKEL_LEFT_FOOT,      frame_id, "left_foot", tm);

      publishTransformTF(user, XN_SKEL_RIGHT_HIP,      frame_id, "right_hip", tm);
      publishTransformTF(user, XN_SKEL_RIGHT_KNEE,     frame_id, "right_knee", tm);
      publishTransformTF(user, XN_SKEL_RIGHT_FOOT,     frame_id, "right_foot", tm);
    }
  }

  void setPose(XnUserID const& user, XnSkeletonJoint const& joint, geometry_msgs::Pose &pose) {
    XnSkeletonJointPosition joint_position;
    g_UserGenerator.GetSkeletonCap().GetSkeletonJointPosition(user, joint, joint_position);
    double x = -joint_position.position.X / 1000.0;
    double y = joint_position.position.Y / 1000.0;
    double z = joint_position.position.Z / 1000.0;

    XnSkeletonJointOrientation joint_orientation;
    g_UserGenerator.GetSkeletonCap().GetSkeletonJointOrientation(user, joint, joint_orientation);

    XnFloat* m = joint_orientation.orientation.elements;

    KDL::Rotation rotation(m[0], m[1], m[2],
                           m[3], m[4], m[5],
                           m[6], m[7], m[8]);
    double qx, qy, qz, qw;
    rotation.GetQuaternion(qx, qy, qz, qw);

    tf::Transform transform;
    transform.setOrigin(tf::Vector3(x, y, z));
    transform.setRotation(tf::Quaternion(qx, -qy, -qz, qw));

    tf::Transform change_frame;
    change_frame.setOrigin(tf::Vector3(0, 0, 0));
    tf::Quaternion frame_rotation;
    //frame_rotation.setEulerZYX(1.5707963268, 0, 1.5707963268);
    //frame_rotation.setEulerZYX(-1.5707963268, 1.5707963268, 0); // camera_frame
    //frame_rotation.setEulerZYX(0, -1.5707963268, -1.5707963268); // camera_frame
    frame_rotation.setEulerZYX(3.14159265359, 0, 0); // camera_frame
    change_frame.setRotation(frame_rotation);

    transform = change_frame * transform;

    pose.position.x = transform.getOrigin().x();
    pose.position.y = transform.getOrigin().y();
    pose.position.z = transform.getOrigin().z();

    tf::Quaternion qq = transform.getRotation();
    pose.orientation.x = qq.x();
    pose.orientation.y = qq.y();
    pose.orientation.z = qq.z();
    pose.orientation.w = qq.w();
  }

  void publishTransforms(const std::string& frame_id, const ros::Time &tm) {
    XnUserID users[15];
    XnUInt16 users_count = 15;
    g_UserGenerator.GetUsers(users, users_count);

    if (users_count < 1) {
      return;
    }

    human_tracking::BodyArray ret;
    ret.header.frame_id = frame_id;
    ret.header.stamp = tm;
    ret.array.resize(users_count);

    //
    for (int i = 0; i < users_count; ++i) {
      XnUserID user = users[i];
      human_tracking::BodyCoords *bd = & (ret.array[i]);
      if (!g_UserGenerator.GetSkeletonCap().IsTracking(user))
        continue;
      bd->link_list.resize(15);
      bd->coords_list.resize(15);
      bd->id = user;

      setPose(user, XN_SKEL_HEAD,          bd->coords_list[0]);
      setPose(user, XN_SKEL_NECK,          bd->coords_list[1]);
      setPose(user, XN_SKEL_TORSO,         bd->coords_list[2]);

      setPose(user, XN_SKEL_LEFT_SHOULDER, bd->coords_list[3]);
      setPose(user, XN_SKEL_LEFT_ELBOW,    bd->coords_list[4]);
      setPose(user, XN_SKEL_LEFT_HAND,     bd->coords_list[5]);

      setPose(user, XN_SKEL_RIGHT_SHOULDER,bd->coords_list[6]);
      setPose(user, XN_SKEL_RIGHT_ELBOW,   bd->coords_list[7]);
      setPose(user, XN_SKEL_RIGHT_HAND,    bd->coords_list[8]);

      setPose(user, XN_SKEL_LEFT_HIP,      bd->coords_list[9]);
      setPose(user, XN_SKEL_LEFT_KNEE,     bd->coords_list[10]);
      setPose(user, XN_SKEL_LEFT_FOOT,     bd->coords_list[11]);

      setPose(user, XN_SKEL_RIGHT_HIP,     bd->coords_list[12]);
      setPose(user, XN_SKEL_RIGHT_KNEE,    bd->coords_list[13]);
      setPose(user, XN_SKEL_RIGHT_FOOT,    bd->coords_list[14]);

      bd->link_list[0] = "head";
      bd->link_list[1] = "neck";
      bd->link_list[2] = "torso";

      bd->link_list[3] = "left_shoulder";
      bd->link_list[4] = "left_elbow";
      bd->link_list[5] = "left_hand";

      bd->link_list[6] = "right_shoulder";
      bd->link_list[7] = "right_elbow";
      bd->link_list[8] = "right_hand";

      bd->link_list[9] = "left_hip";
      bd->link_list[10] = "left_knee";
      bd->link_list[11] = "left_foot";

      bd->link_list[12] = "right_hip";
      bd->link_list[13] = "right_knee";
      bd->link_list[14] = "right_foot";
    }
    pub_bodies_.publish(ret);
  }

  void publishPoints (const std::string &frame_id, const ros::Time &tm) {
    // xn::DepthMetaData depthMD;
    //g_DepthGenerator.GetMetaData(depthMD);
    XnUserID users[16] = {0};
    XnUInt16 users_count = 16;
    g_UserGenerator.GetUsers(users, users_count);

    if(users_count < 1) {
      return;
    }

    xn::SceneMetaData sMD;
    g_UserGenerator.GetUserPixels(0, sMD);
    int width = sMD.XRes();
    int height = sMD.YRes();

    human_tracking::IndicesArray ret;
    ret.header.frame_id = frame_id;
    ret.header.stamp = tm;
    ret.array.resize(users_count);
    ret.id_list.resize(users_count);

    for(int k = 0; k < users_count; k++) {
      int uid = users[k];//
      ret.id_list[k] = uid;
      pcl::PointIndices *idlst = &(ret.array[k]);
      idlst->indices.resize(0);
      const XnLabel* pLabels = sMD.Data();

#if 0 // mirror??
      for(int y = 0; y < height; y++) {
        for(int x = 1; x <= width ; x++) {
          XnLabel label = *pLabels;
          if (label == uid) {
            idlst->indices.push_back(y*width + (width-x));
          }
          //uids[*pLabels]++;
          pLabels++;
        }
      }
#endif
      for(int i = 0; i < width*height; i++) {
        XnLabel label = *pLabels;
        if (label == uid) {
          idlst->indices.push_back(i);
        }
        pLabels++;
      }
    }

    pub_indices_.publish(ret);
  }

  int proc () {
    g_UserGenerator.GetSkeletonCap().SetSkeletonProfile(XN_SKEL_PROFILE_ALL);

    XnStatus nRetVal = g_Context.StartGeneratingAll();
    CHECK_RC(nRetVal, "StartGenerating");

    ros::Rate r(30);

    while (ros::ok()) {
      g_Context.WaitAndUpdateAll();
      ros::Time tm = ros::Time::now();
      publishTransforms(frame_id_, tm);
      publishPoints(frame_id_, tm);
      r.sleep();
    }

    g_Context.Shutdown();
  }

  void SaveCalibration() {
    if ( save_mode_ ) {
      XnUserID aUserIDs[16] = {0};
      XnUInt16 nUsers = 16;
      g_UserGenerator.GetUsers(aUserIDs, nUsers);
      for (int i = 0; i < nUsers; ++i)  {
        // Find a user who is already calibrated
        if (g_UserGenerator.GetSkeletonCap().IsCalibrated(aUserIDs[i])) {
          // Save user's calibration to file
          XnStatus rc = g_UserGenerator.GetSkeletonCap().
            SaveCalibrationDataToFile(aUserIDs[i], calib_fname.c_str());

          ROS_INFO("Cabliration file of User %d is stored to %s -> %d", aUserIDs[i], calib_fname.c_str(), rc);
          break;
        }
      }
    }
  }

  void LoadCalibration(std::string fname) {
    XnUserID aUserIDs[16] = {0};
    XnUInt16 nUsers = 16;
    g_UserGenerator.GetUsers(aUserIDs, nUsers);

    for (int i = 0; i < nUsers; ++i) {
      // Find a user who isn't calibrated or currently in pose
      if (g_UserGenerator.GetSkeletonCap().IsCalibrated(aUserIDs[i])) continue;
      if (g_UserGenerator.GetSkeletonCap().IsCalibrating(aUserIDs[i])) continue;
      // Load user's calibration from file
      XnStatus rc = g_UserGenerator.GetSkeletonCap().
        LoadCalibrationDataFromFile(aUserIDs[i], fname.c_str());

      if (rc == XN_STATUS_OK) {
        // Make sure state is coherent
        g_UserGenerator.GetPoseDetectionCap().StopPoseDetection(aUserIDs[i]);
        g_UserGenerator.GetSkeletonCap().StartTracking(aUserIDs[i]);
      }
      // break;
    }
  }
};

UserTracker *objUserTracker;
XnBool g_bNeedPose   = FALSE;
XnChar g_strPose[20] = "";

void XN_CALLBACK_TYPE User_NewUser(xn::UserGenerator& generator, XnUserID nId, void* pCookie) {
  ROS_INFO("New User %d", nId);

  if ( !objUserTracker->use_calibfile_ ) {
    if (g_bNeedPose) {
      objUserTracker->g_UserGenerator.GetPoseDetectionCap().StartPoseDetection(g_strPose, nId);
    } else {
      objUserTracker->g_UserGenerator.GetSkeletonCap().RequestCalibration(nId, TRUE);
    }
  } else {
#if 1
    XnStatus rc = objUserTracker->g_UserGenerator.GetSkeletonCap().LoadCalibrationDataFromFile(nId, objUserTracker->calib_fname.c_str());

    if (rc == XN_STATUS_OK) {
      ROS_INFO("Use Calib File OK : %s", objUserTracker->calib_fname.c_str());
      // Make sure state is coherent
      objUserTracker->g_UserGenerator.GetPoseDetectionCap().
        StopPoseDetection(nId);
      objUserTracker->g_UserGenerator.GetSkeletonCap().
        StartTracking(nId);
    } else {
      ROS_WARN("Use Calib File Fail ! : %s", objUserTracker->calib_fname.c_str());
    }
#endif
  }
}

void XN_CALLBACK_TYPE User_LostUser(xn::UserGenerator& generator, XnUserID nId, void* pCookie) {
  ROS_INFO("Lost user %d", nId);
}

void XN_CALLBACK_TYPE UserCalibration_CalibrationStart(xn::SkeletonCapability& capability, XnUserID nId, void* pCookie) {
  ROS_INFO("Calibration started for user %d", nId);
}

void XN_CALLBACK_TYPE UserCalibration_CalibrationEnd(xn::SkeletonCapability& capability, XnUserID nId, XnBool bSuccess, void* pCookie) {
  if (bSuccess) {
    ROS_INFO("Calibration complete, start tracking user %d", nId);
    objUserTracker->g_UserGenerator.GetSkeletonCap().StartTracking(nId);
    //
    objUserTracker->SaveCalibration();
  }
  else {
    ROS_INFO("Calibration failed for user %d", nId);
    if (g_bNeedPose) {
      objUserTracker->g_UserGenerator.GetPoseDetectionCap().StartPoseDetection(g_strPose, nId);
    } else {
      objUserTracker->g_UserGenerator.GetSkeletonCap().RequestCalibration(nId, TRUE);
    }
  }
}

void XN_CALLBACK_TYPE UserPose_PoseDetected(xn::PoseDetectionCapability& capability, XnChar const* strPose, XnUserID nId, void* pCookie) {
    ROS_INFO("Pose %s detected for user %d", strPose, nId);
    objUserTracker->g_UserGenerator.GetPoseDetectionCap().StopPoseDetection(nId);
    objUserTracker->g_UserGenerator.GetSkeletonCap().RequestCalibration(nId, TRUE);
}

int main(int argc, char **argv) {
  ros::init(argc, argv, "openni_tracker");

  objUserTracker = new UserTracker();

  XnCallbackHandle hUserCallbacks;
  objUserTracker->g_UserGenerator.RegisterUserCallbacks(User_NewUser,
                                                        User_LostUser,
                                                        NULL, hUserCallbacks);

  if ( !objUserTracker->use_calibfile_ ) {
#if 1
    XnCallbackHandle hCalibrationCallbacks;
    objUserTracker->g_UserGenerator.GetSkeletonCap().
      RegisterCalibrationCallbacks(UserCalibration_CalibrationStart,
                                   UserCalibration_CalibrationEnd,
                                   NULL, hCalibrationCallbacks);

    if (objUserTracker->g_UserGenerator.GetSkeletonCap().NeedPoseForCalibration()) {
      g_bNeedPose = TRUE;
      if (!(objUserTracker->g_UserGenerator.IsCapabilitySupported(XN_CAPABILITY_POSE_DETECTION))) {
        ROS_INFO("Pose required, but not supported");
        return 1;
      }
      XnCallbackHandle hPoseCallbacks;
      objUserTracker->g_UserGenerator.GetPoseDetectionCap().
        RegisterToPoseCallbacks(UserPose_PoseDetected, NULL, NULL, hPoseCallbacks);
      objUserTracker->g_UserGenerator.GetSkeletonCap().GetCalibrationPose(g_strPose);
    }
#endif // for calibration
  }

  return objUserTracker->proc();
  //return 0;
}
